Support co-partitioned range inner equi joins#23184
Conversation
github-actions
Bot
added
physical-expr
gene-bordegaray
changed the title
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
| /// This is optimizer policy: partitioned joins require children that can be | ||
| /// paired by partition index. Inner hash joins can reuse compatible range | ||
| /// partitioning; otherwise the existing hash repartitioning policy applies. | ||
| fn partitioned_join_distribution( |
| /// left: Range(left.a ASC, [10, 20]), required KeyPartitioned(left.a) | ||
| /// right: Range(right.x ASC, [15, 20]), required KeyPartitioned(right.x) | ||
| /// ``` | ||
| pub fn co_partitioned_with( |
There was a problem hiding this comment.
I realized the concept was a bit off. I think this is ok if we haven't had a release... 😅
I linked the mirroring concepts that are used in trino in spark in the PR description
Lesson learned to have consumer of the public API first
There was a problem hiding this comment.
👍 I don't think it's a big deal, this is anyways not the typical method external consumers rely one.
There was a problem hiding this comment.
Without taking a look at the usages of this method, I do notice that the signature seems a bit off: why would it need the required Distribution for checking if two Partitionings are co-partitioned? it seems unrelated.
I'll keep reading though
There was a problem hiding this comment.
Ya I thought similarly but this matters for multi-children operators. For joins, each side has a different may have a diff schema:
left: Partitioning::Range(left.a)
right: Partitioning::Range(right.x)
join: left.a = right.x
So the expressions left.a and right.x are not equal since they are from different children. Because of this we need to check each sides req Distribution to know what each partitioning has to satisfy:
left requirement: key partitioned on left.a
right requirement: key partitioned on right.x
So we need to check:
- left partitioning satisfies left requirement
- right partitioning satisfies right requirement
- the two partition mappings can be paired
I could do something like this:
struct RequiredPartitioning<'a> {
partitioning: &'a Partitioning,
requirement: &'a Distribution,
eq_properties: &'a EquivalenceProperties,
}
impl RequiredPartitioning<'_> {
fn co_partitioned_with(&self, other:
&RequiredPartitioning<'_>) -> bool
}or derive the EquivalenceProperties in the method if we pass the children directly
|
cc: @gabotechs @stuhood |
gene-bordegaray
changed the title
gene-bordegaray
force-pushed
the
gene.bordegaray/2026/06/range-partitioned-joins
branch
from
8583303 to
45d598b
Compare
| # TEST 3b: Non-Range Join Repartitions | ||
| # The source Range partitioning does not satisfy a join on non_range_key, so | ||
| # planning still inserts Hash repartitioning. | ||
| ########## |
There was a problem hiding this comment.
I can propbably eliminate this, almost a sanity check / nice to see after teh positive case
|
Amazing timing. Will get you some feedback on this by early next week! Thank you! |
@stuhood great, thank you! I also have a huge line of follow up issues to support more join types that we can split up 👍 I am trying to make smaller tickets as more of the plumbing gets in to get more people involved |
gene-bordegaray
changed the title
gene-bordegaray
force-pushed
the
gene.bordegaray/2026/06/range-partitioned-joins
branch
from
45d598b to
2fca2bb
Compare
|
I have added a stacked PR that shows how I imagine Aggregations will consume the |
gabotechs
left a comment
gabotechs
left a comment
There was a problem hiding this comment.
This is looking good! left a first round of comments.
The main one is about trying to collapse Distribution::KeyPartitioned and Distribution::HashPartitioned into just a single Distribution variant, I do see that in pretty much all places the code that handles both is the same, and I get the feeling that Distribution::HashPartitioned is just an unfortunate name that was choosen a while back, but it's not really coupled to hashes.
| /// Requires rows with equal values for the given keys to be colocated in | ||
| /// the same partition, without requiring a specific partitioning algorithm. | ||
| /// | ||
| /// Unlike [`Self::HashPartitioned`], this can be satisfied by non-hash | ||
| /// partitioning such as range partitioning. A partitioning on a subset of | ||
| /// these keys can also satisfy this requirement because rows equal on all | ||
| /// required keys are also equal on any subset. | ||
| /// | ||
| /// For multi-input operators, satisfaction alone is not enough: each input | ||
| /// may satisfy its own key requirement while using incompatible partition | ||
| /// boundaries. Use [`Partitioning::co_partitioned_with`] before pairing | ||
| /// partitions by index. | ||
| KeyPartitioned(Vec<Arc<dyn PhysicalExpr>>), | ||
| } |
There was a problem hiding this comment.
I'm wondering if there's a way of condense the two HashPartitioned and KeyPartitioned into just one. Take a look at the Distribution::HashPartitioned docs:
Requires children to be distributed in such a way that the same values of the keys end up in the same partition
For meeting that distribution, I don't think it's really necessary that children are partitioned specifically with a hashin algorithm, they can be partitioned in any way that satisfies that invariant (same value of the keys end up in the same partition).
I'm yet to read the full PR, but I'm getting the feeling that the previous HashPartitioned naming is just unfortunate, and could probably have been just named KeyPartitioned from the beginning.
What's your opinion on trying to rename HashPartitioned to KeyPartitioned instead?
There was a problem hiding this comment.
ya I agree with this long-term 🤔
What do you think about deprecating HashPartitioned then Introducing KeyPartitioned and lettting this PR be the first real consumer of it. Then we don't allow range to satisfy KeyPartitioned generally here, and only allow range to satisfy the distribution in the partitioend join case.
Then in following PRs we slowly add some provate helpers to allow range to satisfay the KeyPartitioned like aggregations, then once its at a reasonable scope (a good amount of operators are working and tested thoroughly), we and range to the general satisfaction() public method.
Trying to think how to do this incrementally
| /// left: Range(left.a ASC, [10, 20]), required KeyPartitioned(left.a) | ||
| /// right: Range(right.x ASC, [15, 20]), required KeyPartitioned(right.x) | ||
| /// ``` | ||
| pub fn co_partitioned_with( |
There was a problem hiding this comment.
👍 I don't think it's a big deal, this is anyways not the typical method external consumers rely one.
| /// left: Range(left.a ASC, [10, 20]), required KeyPartitioned(left.a) | ||
| /// right: Range(right.x ASC, [15, 20]), required KeyPartitioned(right.x) | ||
| /// ``` | ||
| pub fn co_partitioned_with( |
There was a problem hiding this comment.
Without taking a look at the usages of this method, I do notice that the signature seems a bit off: why would it need the required Distribution for checking if two Partitionings are co-partitioned? it seems unrelated.
I'll keep reading though
| } | ||
| equivalent_exprs(left_exprs, right_exprs, eq_properties) | ||
| } | ||
| (Partitioning::Hash(_, _), Partitioning::Hash(_, _)) => true, |
There was a problem hiding this comment.
🤔 Is this correct?
What if the left Partitioning::Hash contains 1 expression but the right Partitioning::Hash contains 6 completely different expressions? I'm not sure if returning true in that case is correct.
There was a problem hiding this comment.
This is only correct check before:
if !self
.satisfaction(required, eq_properties, false)
.is_satisfied()
|| !other
.satisfaction(other_required, other_eq_properties, false)
.is_satisfied()
{
return false;
}since this makes sure each side satisfies its join distribution. So Hash(a) vs Hash(b, c, ...) wouldnt reach this.
Something like this may up front be clenaer and more obvious though:
let left_satisfies = matches!(
self.satisfaction(required,
eq_properties, false),
PartitioningSatisfaction::Exact
);
let right_satisfies = matches!(
other.satisfaction(other_required,
other_eq_properties, false),
PartitioningSatisfaction::Exact
);
if !left_satisfies || !right_satisfies {
return false;
}| Distribution::HashPartitioned(required_exprs) => match self { | ||
| // Here we do not check the partition count for hash partitioning and assumes the partition count | ||
| // and hash functions in the system are the same. In future if we plan to support storage partition-wise joins, | ||
| // then we need to have the partition count and hash functions validation. | ||
| Partitioning::Hash(partition_exprs, _) => { | ||
| // Empty hash partitioning is invalid | ||
| if partition_exprs.is_empty() || required_exprs.is_empty() { | ||
| return PartitioningSatisfaction::NotSatisfied; | ||
| } | ||
|
|
||
| if equivalent_exprs(required_exprs, partition_exprs, eq_properties) { | ||
| return PartitioningSatisfaction::Exact; | ||
| } | ||
|
|
||
| let eq_groups = eq_properties.eq_group(); | ||
| if !eq_groups.is_empty() { | ||
| if allow_subset { | ||
| let normalized_partition_exprs = | ||
| normalize_exprs(partition_exprs, eq_properties); | ||
| let normalized_required_exprs = | ||
| normalize_exprs(required_exprs, eq_properties); | ||
| if Self::is_subset_partitioning( | ||
| &normalized_partition_exprs, | ||
| &normalized_required_exprs, | ||
| ) { | ||
| return PartitioningSatisfaction::Subset; | ||
| } | ||
| } | ||
| } else if allow_subset | ||
| && Self::is_subset_partitioning(partition_exprs, required_exprs) | ||
| { | ||
| return PartitioningSatisfaction::Subset; | ||
| } | ||
|
|
||
| PartitioningSatisfaction::NotSatisfied | ||
| } | ||
| Partitioning::Hash(partition_exprs, _) => Self::key_expr_satisfaction( | ||
| partition_exprs, | ||
| required_exprs, | ||
| eq_properties, | ||
| allow_subset, | ||
| ), | ||
| Partitioning::RoundRobinBatch(_) |
There was a problem hiding this comment.
The fact that satisfying both Distribution::KeyPartitioned and Distribution::HashPartitioned with Partitioning::Hash requires exactly the same code reinforces my comment above. The longer I look the more I think we should collapse Distribution::HashPartitioned and Distribution::KeyPartitioned into just 1
| hash_exprs: Vec<Arc<dyn PhysicalExpr>>, | ||
| n_target: usize, | ||
| allow_subset_satisfy_partitioning: bool, | ||
| force_to_target: bool, |
There was a problem hiding this comment.
😬 if-driven-development alert
There was a problem hiding this comment.
I'll keep reading to see if I can suggest something better
There was a problem hiding this comment.
Ok, it's not that bad, but at this point I don't think the add_hash_on_top is working well as an abstraction.
The output I would expect from this function is that it adds a hash repartition on top, but it might not do that, it depends on some conditional logic. There's some relatively good docs that make up for it, but rather than documenting the weirdness, it might be worth not having it at all from the beginning.
If you ask me, I'd just remove this function and inline the body in the only place where it's used:
Code suggestion
Distribution::HashPartitioned(exprs)
| Distribution::KeyPartitioned(exprs) => 'curr_match_branch: {
// See https://github.com/apache/datafusion/issues/18341#issuecomment-3503238325 for background
// When inserting hash is necessary to satisfy hash requirement, insert hash repartition.
if !hash_necessary {
break 'curr_match_branch;
}
// Early return if hash repartition is unnecessary
// `RepartitionExec: partitioning=Hash([...], 1), input_partitions=1` is unnecessary.
if target_partitions == 1 && current_partitions == 1 {
break 'curr_match_branch;
}
let is_satisfied = child
.plan
.output_partitioning()
.satisfaction(
&Distribution::HashPartitioned(exprs.to_vec()),
child.plan.equivalence_properties(),
allow_subset_satisfy_partitioning,
)
.is_satisfied();
// Add hash repartitioning when:
// - When subset satisfaction is enabled (current >= threshold): only repartition if not satisfied
// - When below threshold (current < threshold): repartition if expressions don't match OR to increase parallelism
let needs_repartition = if force_hash_to_target {
!is_satisfied || target_partitions != current_partitions
} else if allow_subset_satisfy_partitioning {
!is_satisfied
} else {
!is_satisfied || target_partitions > current_partitions
};
if needs_repartition {
let part = Partitioning::Hash(exprs.to_vec(), target_partitions);
child.plan = Arc::new(
RepartitionExec::try_new(child.plan, part)?
.with_preserve_order(),
);
}
}Which allows trimming a fair amount of code (35 insertions(+), 81 deletions(-)) and is easier to track as readers do not need to jump back and forth
There was a problem hiding this comment.
This is a nit though, what you have right now in this PR is also good, so only take the suggestion if you want.
There was a problem hiding this comment.
Ya I think you are probably right, I am going to play with this once we talk about the larger issues first and see what fits nicely 👍
|
I think Key idea
Hash partitioning and range partitioning are two concrete ways to satisfy this requirement. More precisely: Equivalently, different output partitions contain disjoint sets of key values. This seems to cover the logical requirement that Example 1: aggregation has one input
That means all rows with the same For example, both of the following satisfy this requirement:
In either case, we do not need to insert an extra hash repartition before the aggregation. Example 2: hash join has two inputsFor a partitioned hash join, it is not enough for each side to be independently The two sides must also be co-partitioned with respect to the join keys. That means equal join-key values from the left and right inputs must be assigned to the same partition id. The physical execution pattern is: So the required co-partitioning rule is: This implies some extra compatibility requirements across the two inputs:
So I think there are two related concepts: Implementation planHere is my initial thoughts on a implementation plan (note I haven't read the related code carefully, so it's just some rough ideas)
|
4 participants
Which issue does this PR close?
Rationale for this change
DataFusion can represent source-declared range partitioning, but partitioned hash joins still required hash partitioned inputs. So an inner join on compatible range-partitioned keys would insert unnecessary hash repartitions, even when each left/right partition already covered the same key domain.
This PR adds a partitioning requirement that means "equal key values are co-located" . I was calling this "compatibility" but found we can satisfy the requirement with looser conditions. Other systems call this "co-location" or "co-partitioning" (trino, spark). Which they (and now I am proposing) define as when both sides of a join are already partitioned so matching key values appear in corresponding partitions, so we can join partition pairs directly without repartitioning the sides.
This lets "co-partitioned" range inputs satisfy inner partitioned hash joins. This will also be applicable to other join types and operators but kept the first PR thin to keep scope more reviewable.
What changes are included in this PR?
Adds
Distribution::KeyPartitioned(Vec<Arc<dyn PhysicalExpr>>)as a public distribution requirement.HashPartitioned([a])means rows must be partitioned by hash ona.KeyPartitioned([a])means rows with equalavalues must be co-located, but the partitioning algorithm may be hash, range, or another compatible scheme.Adds
Partitioning::co_partitioned_with(...)to validate that two independently satisfying partitionings also can be paired by partition index.Changes inner partitioned
HashJoinExecrequirements fromHashPartitionedtoKeyPartitioned.HashPartitionedfor now.Updates
EnforceDistributionso co-partitioned range inner joins avoid repartitioning.idoes not represent the same key domain on both sides.KeyPartitionedin this PR.Keeps partitioned dynamic filter pushdown restricted to hash-compatible routing.
Degrades range join output partitioning to
UnknownPartitioning(n)rather than erroring. Adding this behavior would need more tests and careful thought about, I think its safert o just degrade for first PR.Are these changes tested?
Yes.
KeyPartitionedsatisfaction for hash and range partitioning.co_partitioned_withfor compatible and incompatible range/hash partitioning.EnforceDistributionbehavior for:Are there any user-facing changes?
Yes.
This PR changes public physical planning APIs:
Distribution::KeyPartitioned.Partitioning::co_partitioned_with.Distribution.